#include <iostream>
#include <queue>
#include <vector>
#include <cassert>

using namespace std;

// 215. 数组中的第K个最大元素
// 给定整数数组 nums 和整数 k，请返回数组中第 k 个最大的元素。
// 请注意，你需要找的是数组排序后的第 k 个最大的元素，而不是第 k 个不同的元素。
// 你必须设计并实现时间复杂度为 O(n) 的算法解决此问题。

// 使用快排分区
// 时间复杂度：O(n)
// 空间复杂度： O(1)
class Solution {
private:
    int partition(vector<int> &arr, int l, int r) {
//        srand(NULL);  // 将其写在入口函数中初始化
        int p = l + rand() %(r - l + 1);    // [p, r)
        swap(arr[l], arr[p]);   // pivot放在最左端arr[l]
        // 注意是从大到小的分区，所以前面区间的数大于后面区间的数
        int lt = l+1;  // [l+1...lt) > p ; [lt...i) < p
        for(int i = l + 1; i <= r; i++)
            if(arr[i] > arr[l])
                swap(arr[i], arr[lt++]);

        swap(arr[l], arr[lt-1]);

        return lt-1;
    }
    int findKthLargest(vector<int>& nums, int l, int r, int k) {
        if( l == r)
            return nums[l];

        int p = partition(nums, l, r);

        if(p == k)
            return nums[p];
        else if( k < p)
            return findKthLargest(nums, l, p-1, k);
        else {
            assert(k > p);
            return findKthLargest(nums, p+1, r, k);
        }
    }
public:
    int findKthLargest(vector<int>& nums, int k) {
        assert(nums.size() > 0 && k > 0 && k <= nums.size());
        srand(time(NULL));
        return findKthLargest(nums, 0, nums.size()-1, k-1); //k-1是为了和数组下标一致
    }
};

int main() {
    vector<int> nums1 = {3, 2, 1, 5, 6, 4};
    cout << Solution().findKthLargest(nums1, 2) << endl;
    // 5

    vector<int> nums2 = {3, 1, 2, 4};
    cout << Solution().findKthLargest(nums2, 2) << endl;
    // 3

    vector<int> nums3 = {3, 3, 3, 3, 3, 3, 3, 3, 3};
    cout << Solution().findKthLargest(nums3, 8) << endl;
    // 3
    return 0;
}
